#include "EngineWebGPU.h"
#include <iostream>
#ifdef __EMSCRIPTEN__
#include <emscripten.h>
#include <emscripten/html5.h>
#include <emscripten/html5_webgpu.h>
#else
#include <dawn/dawn_proc.h>
#include <dawn/native/DawnNative.h>
#include <webgpu/webgpu_cpp.h>
#include <thread>
#include <webgpu/webgpu_glfw.h>
#endif


wgpu::TextureFormat WGPUConvertUnormToSRGB(wgpu::TextureFormat Format)
{
  switch (Format)
  {
  case wgpu::TextureFormat::RGBA8Unorm:
    return wgpu::TextureFormat::RGBA8UnormSrgb;
  case wgpu::TextureFormat::BGRA8Unorm:
    return wgpu::TextureFormat::BGRA8UnormSrgb;
  default:
    return Format;
  }
}

void InstancePollEvents(WGPUInstance wgpuInstance) {
  wgpuInstanceProcessEvents(wgpuInstance);
}

#ifndef __EMSCRIPTEN__
GLFWwindow* initGLFW() {
  if (!glfwInit()) {
    std::cerr << "Failed to initialize GLFW" << std::endl;
    return nullptr;
  }

  glfwWindowHint(GLFW_CLIENT_API, GLFW_NO_API);
  GLFWwindow* window = glfwCreateWindow(800, 600, "WebGPU Test", nullptr, nullptr);

  if (!window) {
    glfwTerminate();
    std::cerr << "Failed to create GLFW window" << std::endl;
  }
  return window;
}
#endif

int configureSurface(
  const wgpu::Surface& surface,
  const wgpu::Adapter& wgpuAdapter,
  const wgpu::Device& wgpuDevice,
  uint32_t width,
  uint32_t height,
  wgpu::TextureFormat* wgpuRTVFormats
) {
#ifndef __EMSCRIPTEN__
  wgpu::SurfaceCapabilities wgpuSurfaceCapabilities{};
  auto status = surface.GetCapabilities(wgpuAdapter, &wgpuSurfaceCapabilities);
  if (status != wgpu::Status::Success) {
    std::cerr << "Failed to get surface capabilities" << std::endl;
    return -1;
  }

  if (wgpuSurfaceCapabilities.formatCount == 0) {
    std::cerr << "No supported surface formats" << std::endl;
    return -1;
  }

  wgpu::TextureFormat wgpuPreferredFormat = wgpuSurfaceCapabilities.formats[0];
#else
  wgpu::TextureFormat wgpuPreferredFormat = wgpu::TextureFormat::BGRA8Unorm;
#endif

  // WGPUTextureFormat
  wgpuRTVFormats[0] = wgpuPreferredFormat;
  wgpuRTVFormats[1] = wgpu::TextureFormat::Undefined;

#ifndef __EMSCRIPTEN__
  for (uint32_t i = 0; i < wgpuSurfaceCapabilities.formatCount; ++i) {
    if (wgpuSurfaceCapabilities.formats[i] == WGPUConvertUnormToSRGB(wgpuPreferredFormat)) {
      wgpuRTVFormats[1] = wgpuSurfaceCapabilities.formats[i];
      break;
    }
  }
#endif

  wgpu::SurfaceConfiguration wgpuSurfaceConfig{};
  wgpuSurfaceConfig.viewFormatCount = 2;
  wgpuSurfaceConfig.nextInChain = nullptr;
  wgpuSurfaceConfig.device = wgpuDevice;
  wgpuSurfaceConfig.usage = wgpu::TextureUsage::RenderAttachment;
  wgpuSurfaceConfig.format = wgpuPreferredFormat;
  wgpuSurfaceConfig.width = static_cast<uint32_t>(width);
  wgpuSurfaceConfig.height = static_cast<uint32_t>(height);
  wgpuSurfaceConfig.presentMode = wgpu::PresentMode::Fifo; // 回退到 FIFO

#ifndef __EMSCRIPTEN__
  for (uint32_t i = 0; i < wgpuSurfaceCapabilities.presentModeCount; ++i) {
    if (wgpuSurfaceCapabilities.presentModes[i] == wgpu::PresentMode::Mailbox) {
      wgpuSurfaceConfig.presentMode = wgpu::PresentMode::Mailbox;
      break;
    }
  }
#endif
  wgpuSurfaceConfig.alphaMode = wgpu::CompositeAlphaMode::Opaque;
  wgpuSurfaceConfig.viewFormats = wgpuRTVFormats;
  wgpuSurfaceConfig.viewFormatCount = (wgpuRTVFormats[1] != wgpu::TextureFormat::Undefined) ? 2 : 1;

  // 应用表面配置
  surface.Configure(&wgpuSurfaceConfig);

  return 0;
}

EngineWebGPU::EngineWebGPU() {
#ifndef __EMSCRIPTEN__
  dawnProcSetProcs(&dawn::native::GetProcs());
#endif
}

EngineWebGPU::~EngineWebGPU() {
#ifndef __EMSCRIPTEN__
  if (mWindow) {
    glfwDestroyWindow(mWindow);
    glfwTerminate();
  }
#endif
}

void EngineWebGPU::init() {
#ifndef __EMSCRIPTEN__
  mWindow = initGLFW();
#endif

  createInstance();
  findAdapter();
  createDevice();
  createSurface();

  configureSurface(mSurface, mAdapter, mDevice, 800, 600, wgpuRTVFormats);

  mQueue = mDevice.GetQueue();
}

void EngineWebGPU::createInstance() {
#ifndef __EMSCRIPTEN__
  wgpu::InstanceDescriptor instanceDesc{};
  instanceDesc.nextInChain = nullptr;
  mInstance = wgpu::CreateInstance(&instanceDesc);
#else
  mInstance = wgpu::CreateInstance(nullptr);
#endif

  if (!mInstance) {
    throw std::runtime_error("Failed to create WGPU instance");
  }
}

void EngineWebGPU::findAdapter() {
#ifndef __EMSCRIPTEN__
  wgpu::RequestAdapterOptions requestOptions;
  requestOptions.forceFallbackAdapter = false;
  requestOptions.backendType = wgpu::BackendType::Metal;

  AsyncResult<wgpu::Adapter> adapterResult;

  mInstance.RequestAdapter(
    &requestOptions,
    wgpu::CallbackMode::AllowProcessEvents,
    [&adapterResult](wgpu::RequestAdapterStatus status, wgpu::Adapter adapter, wgpu::StringView message) {
      adapterResult.status = status;
      adapterResult.result = adapter;
      adapterResult.message = message;
    });

  while (!adapterResult.result) {
    InstancePollEvents(mInstance.Get());
  }
  mAdapter = adapterResult.result;
#endif
}


void EngineWebGPU::createDevice() {
#ifdef __EMSCRIPTEN__
  double css_w = 0.0;
  double css_h = 0.0;

  int win_width = 1280;
  int win_height = 720;
  float dpi = 1.f;

  emscripten_get_element_css_size("canvas", &css_w, &css_h);
  emscripten_get_canvas_element_size("canvas", &win_width, &win_height);
  dpi = emscripten_get_device_pixel_ratio();
  win_width = static_cast<int>(std::round(css_w * dpi));
  win_height = static_cast<int>(std::round(css_h * dpi));

  emscripten_set_canvas_element_size("canvas", win_width, win_height);
  std::cout << "Canvas size: " << win_width << "x" << win_height << "\n";
  mDevice = wgpu::Device::Acquire(emscripten_webgpu_get_device());
  if (!mDevice) {
    std::cerr << "Failed to acquire WebGPU device\n";
  }
  mDevice.SetUncapturedErrorCallback(
    [](WGPUErrorType type, const char* msg, void* data) {
      std::cout << "WebGPU Error: " << msg << "\n" << std::flush;
    },
    nullptr);
#else
  std::vector<wgpu::FeatureName> requiredFeatures;  // empty for now

  wgpu::DeviceDescriptor deviceDesc;
  deviceDesc.requiredFeatureCount = requiredFeatures.size();
  deviceDesc.requiredFeatures = requiredFeatures.data();
  deviceDesc.SetUncapturedErrorCallback(
    [](const wgpu::Device& device, wgpu::ErrorType type, wgpu::StringView message) {
      std::cout << "Error: " << message.data << std::endl;
    });

  mDevice = mAdapter.CreateDevice(&deviceDesc);
#endif
}

void EngineWebGPU::createSurface()
{
  wgpu::SurfaceDescriptor descriptor;
#ifdef __EMSCRIPTEN__
  wgpu::SurfaceDescriptorFromCanvasHTMLSelector selector;
  selector.selector = "canvas";
  descriptor.nextInChain = &selector;
#else
  auto chainedDescriptor = wgpu::glfw::SetupWindowAndGetSurfaceDescriptor(mWindow);
  descriptor.nextInChain = chainedDescriptor.get();
#endif

  mSurface = mInstance.CreateSurface(&descriptor);
  if (!mSurface) {
    std::cerr << "Failed to create surface" << std::endl;
  }
}

wgpu::ShaderModule EngineWebGPU::createShaderModule(const char* shaderSource) {
  wgpu::ShaderModuleWGSLDescriptor wgslDesc{};
  wgslDesc.code = shaderSource;

  wgpu::ShaderModuleDescriptor shaderDesc{};
  shaderDesc.nextInChain = &wgslDesc;

  wgpu::ShaderModule shaderModule = mDevice.CreateShaderModule(&shaderDesc);
  if (!shaderModule) {
    std::cerr << "Failed to create shader module" << std::endl;
  }
  return std::move(shaderModule);
}

wgpu::PipelineLayout EngineWebGPU::createPipelineLayout() {
  // 创建管线布局（这里不需要绑定组，所以使用空的布局）
  wgpu::PipelineLayoutDescriptor pipelineLayoutDesc{};
  pipelineLayoutDesc.bindGroupLayoutCount = 0;
  pipelineLayoutDesc.bindGroupLayouts = nullptr;
  wgpu::PipelineLayout pipelineLayout = mDevice.CreatePipelineLayout(&pipelineLayoutDesc);
  if (!pipelineLayout) {
    std::cerr << "Failed to create pipeline layout" << std::endl;
  }
  return std::move(pipelineLayout);
}

wgpu::RenderPipeline EngineWebGPU::createRenderPipeline(
  wgpu::PipelineLayout pipelineLayout,
  wgpu::ShaderModule shaderModule
) {
  // 设置渲染管线
  wgpu::RenderPipelineDescriptor pipelineDesc{};
  pipelineDesc.layout = pipelineLayout;

  // 顶点状态（无顶点缓冲区）
  pipelineDesc.vertex.module = shaderModule;
  pipelineDesc.vertex.entryPoint = "vs_main";
  pipelineDesc.vertex.bufferCount = 0;
  pipelineDesc.vertex.buffers = nullptr;

  // 片段状态
  wgpu::ColorTargetState colorTarget{};
  colorTarget.format = wgpuRTVFormats[0];
  colorTarget.blend = nullptr;
  colorTarget.writeMask = wgpu::ColorWriteMask::All; // 默认混合状态

  wgpu::FragmentState fragmentState{};
  fragmentState.module = shaderModule;
  fragmentState.entryPoint = "fs_main";
  fragmentState.targetCount = 1;
  fragmentState.targets = &colorTarget;

  // 其他管线设置
  pipelineDesc.label = "SwapChainPresentPSO";
  pipelineDesc.primitive.topology = wgpu::PrimitiveTopology::TriangleList;
  pipelineDesc.primitive.stripIndexFormat = wgpu::IndexFormat::Undefined;
  pipelineDesc.primitive.frontFace = wgpu::FrontFace::CCW;
  pipelineDesc.primitive.cullMode = wgpu::CullMode::None;
  pipelineDesc.multisample.count = 1;
  pipelineDesc.fragment = &fragmentState;
  pipelineDesc.multisample.mask = 0xFFFFFFFF;
  pipelineDesc.multisample.alphaToCoverageEnabled = false;

  wgpu::RenderPipeline renderPipeline = mDevice.CreateRenderPipeline(&pipelineDesc);
  if (!renderPipeline) {
    std::cerr << "Failed to create render pipeline" << std::endl;
  }
  return std::move(renderPipeline);
}

wgpu::TextureView EngineWebGPU::getSuffaceTextureView() {
  wgpu::SurfaceTexture surfaceTexture;
  mSurface.GetCurrentTexture(&surfaceTexture);
#ifdef __EMSCRIPTEN__
  if (surfaceTexture.status != wgpu::SurfaceGetCurrentTextureStatus::Success)
  {
    std::cout << "wgpu::Surface::GetCurrentTexture failed: " << std::endl;
    return nullptr;
  }
#else
  if (surfaceTexture.status != wgpu::SurfaceGetCurrentTextureStatus::SuccessOptimal)
  {
    std::cout << "wgpu::Surface::GetCurrentTexture failed: " << std::endl;
    return nullptr;
  }
#endif

  wgpu::Texture wgpuTexture = surfaceTexture.texture;
  wgpu::TextureViewDescriptor textureViewDesc{};
  textureViewDesc.format = wgpuRTVFormats[0]; // Use sRGB format if needed
  textureViewDesc.label = "SwapChainPresentTextureView";
  textureViewDesc.dimension = wgpu::TextureViewDimension::e2D;
  textureViewDesc.mipLevelCount = 1;
  textureViewDesc.arrayLayerCount = 1;
  textureViewDesc.aspect = wgpu::TextureAspect::All;

  wgpu::TextureView textureView = wgpuTexture.CreateView(&textureViewDesc);
  return std::move(textureView);
}

void EngineWebGPU::loopPC(wgpu::RenderPipeline renderPipeline) {
  #ifndef __EMSCRIPTEN__
  while (!glfwWindowShouldClose(mWindow)) {
    glfwPollEvents();
    InstancePollEvents(mInstance.Get());
    loop(renderPipeline);

    // 呈现表面
    mSurface.Present();
  }
  #endif
}

void EngineWebGPU::loop(wgpu::RenderPipeline renderPipeline) {
  wgpu::CommandEncoderDescriptor encoderDesc{};
  wgpu::CommandEncoder encoder = mDevice.CreateCommandEncoder(&encoderDesc);

  // 设置渲染通道
  wgpu::RenderPassColorAttachment colorAttachment{};
  colorAttachment.view = getSuffaceTextureView();
  colorAttachment.loadOp = wgpu::LoadOp::Clear;
  colorAttachment.storeOp = wgpu::StoreOp::Store;
  colorAttachment.clearValue = { 0.0, 0.0, 0.0, 0.0 }; // 黑色背景
  colorAttachment.depthSlice = WGPU_DEPTH_SLICE_UNDEFINED;

  wgpu::RenderPassDescriptor renderPassDesc{};
  renderPassDesc.colorAttachmentCount = 1;
  renderPassDesc.colorAttachments = &colorAttachment;
  renderPassDesc.depthStencilAttachment = nullptr;

  wgpu::RenderPassEncoder renderPass = encoder.BeginRenderPass(&renderPassDesc);
  if (!renderPass) {
    std::cerr << "Failed to begin render pass" << std::endl;
    return;
  }

  int width = 800;
  int height = 600;
  renderPass.SetViewport(0, 0, width, height, 0.0f, 1.0f);
  renderPass.SetScissorRect(0, 0, width, height);
  renderPass.SetPipeline(renderPipeline.Get());
  renderPass.Draw(3, 1, 0, 0);
  renderPass.End();

  // 提交命令
  wgpu::CommandBufferDescriptor cmdBufferDesc{};
  wgpu::CommandBuffer commandBuffer = encoder.Finish(&cmdBufferDesc);
  mQueue.Submit(1, &commandBuffer);
}
